import os, pydicom, copy, yaml, numpy as np
from PIL import Image
from docx import Document
from docx.oxml.shared import OxmlElement, qn
from docx.shared import Cm
from datetime import datetime
from tkinter import filedialog
from tqdm import tqdm
# import SimpleITK as sitk

class FDATool():
    def __init__(self, dicom_path):
        self.dicom_path = str(dicom_path)
        self.config_file = 'config_info.yml'
        self._parse_config()

    def _parse_config(self):
        with open(self.config_file, 'r', encoding='utf-8') as f:
            config = yaml.safe_load(f)
        self.tag_info = config['tags']
        for tag in self.tag_info.values():
            try:
                tag['id'] = tuple(int(item, base=16) for item in tag['id'].split(','))
            except:
                pass
        self.center_dcm = config['center_dcm']['num']
        self.center_step = config['center_dcm']['step']
        self.include = [i.lower() for i in config['include_body_part']]
        self.exclude = [i.lower() for i in config['exclude_body_part']]

    def get_tag_value(self, dicom_path, tag_info):
        tag_value = copy.deepcopy(tag_info)
        if dicom_path.endswith('.dcm'):
            try:
                df = pydicom.read_file(dicom_path)
            except Exception as err:
                print(err)
                df = {}
            
            for tag in tag_value.values():
                try:
                    tag['value'] = df[tag['id']].value
                except KeyError:
                    tag['value'] = 'N.A.'
        return tag_value

    def convert_jpg(self, dcm_image_path):
        jpg_path = dcm_image_path.replace('dcm', 'jpg')

        try:
            ds = pydicom.read_file(dcm_image_path)  # 读取dicom文件的相关信息
            img_array = ds.pixel_array  # 获取array
        
        except:
            img_array = np.ones([256, 256])
        
        # ds = sitk.ReadImage(dcm_image_path)
        # pixel_array = sitk.GetArrayFromImage(ds)
        # img_array = np.squeeze(pixel_array)
        
        self.__convert_from_dicom_to_jpg(img_array, jpg_path)  # 调用函数，转换成jpg文件并保存到对应的路径
        return jpg_path

    def __convert_from_dicom_to_jpg(self, array, save_path):
        if len(array.shape) == 2:
            img = Image.fromarray(self.__normalize(array))
        else:
            img = Image.fromarray(array)
        img.save(save_path, quality=95)
    
    def __normalize(self, array):
        img = (array - np.min(array)) / (np.max(array) - np.min(array))
        return (img * 255).astype('uint8')

    def extract_parameter(self):
        result_content = {}
        center_dcm = {}
        try:
            for root, dir_path, filenames in os.walk(self.dicom_path):
                if filenames:
                    print('Extracting parameters from %s' % root)
                    protocal, body_part, series_name = root.split(os.sep)[-3:]
                    
                    if self.include != ["all"] and body_part.lower() not in self.include:
                        print('Skip body part:%s, it is not in the include path!' % body_part)
                        continue
                    
                    if body_part.lower() in self.exclude:
                        print('Skip body part:%s, it is in the exclude path!' % body_part)
                        continue
                    
                    if protocal not in result_content:
                        result_content[protocal] = {}
                        center_dcm[protocal] = {}

                    if body_part not in result_content[protocal]:
                        result_content[protocal][body_part] = {}
                        center_dcm[protocal][body_part] = {}
                    
                    temp_dict = {}
                    dcm_num = [i.endswith('.dcm') for i in filenames].count(True)
                    for filename in filenames:
                        if filename.endswith('.dcm'):
                            fullname = os.path.join(root, filename)
                            temp_dict[fullname] = self.get_tag_value(fullname, self.tag_info)
                            try:
                                instance_num = int(temp_dict[fullname]['Instance']['value'])
                            except:
                                instance_num = 0
                            
                            description = []
                            for item, item_info in temp_dict[fullname].items():
                                value = str(item_info['value'])
                                if 'skip' in item_info.keys() and item_info['skip']:
                                    continue
                                if 'stitching' in item_info.keys():
                                    content = value
                                    if 'N.A.' != value:
                                        content += str(temp_dict[fullname][item_info['stitching']]['value'])
                                    append_content = item + ' : ' + content
                                else:
                                    if 'N.A.' == value:
                                        pass
                                    else:
                                        if 'resolution' in item.lower():
                                            value = str(item_info['value']).split("I")[0]
                                        elif 'thickness' in item.lower():
                                            value = '%.1f' % float(value)
                                        
                                    append_content = item + ' : ' + value
                                description.append(append_content)

                            if series_name not in result_content[protocal][body_part]:
                                result_content[protocal][body_part][series_name] = {}
                                center_dcm[protocal][body_part][series_name] = {}

                            result_content[protocal][body_part][series_name][fullname] = '\n'.join(description)

                            if self.__is_center_dcm(dcm_num, instance_num):
                                center_dcm[protocal][body_part][series_name][instance_num] = fullname

        except Exception as err:
            print('In dicom %s encounter an error\n %s' % (fullname, err))

        finally:
            print('Extracting parameters done!')
            return result_content, center_dcm

    def __is_center_dcm(self, total_num, instance_num):
        center_num = (self.center_dcm - 1) * (self.center_step) + 1
        first_num = 1 if total_num <= center_num else int((total_num - center_num) / 2 + 1)
        center_list = range(first_num, first_num + center_num, self.center_step)
        return instance_num in center_list
            
    def write_2_docx(self, content, center_info):
        if os.path.exists('Extracted'):
            pass
        else:
            os.mkdir('Extracted')
        time_append = datetime.strftime(datetime.now(), '%Y-%m-%d_%H%M%S')
        docx_name = 'Extract_Parameter_' + time_append + '.docx'
        docx_path = 'Extracted\\' + docx_name
        print('Write parameters to docx %s' % docx_path)
        docx = Document()
        docx.styles['Normal'].font.name = u'Times New Roman'
        docx.styles['Normal']._element.rPr.rFonts.set(qn('w:eastAsia'), u'Times New Roman')

        for protocal, protocal_value in tqdm(content.items()):
            title_run = docx.add_paragraph(protocal, style="Heading 1").add_run()
            title_run.font.name = u'Times New Roman'
            title_run._element.rPr.rFonts.set(qn('w:eastAsia'), u'Times New Roman')

            for body_part, body_part_value in protocal_value.items():
                title_run = docx.add_paragraph(body_part, style="Heading 2").add_run()
                title_run.font.name = u'Times New Roman'
                title_run._element.rPr.rFonts.set(qn('w:eastAsia'), u'Times New Roman')

                for series_name, series_info in body_part_value.items():
                    table = docx.add_table(rows=1, cols=2, style='Table Grid')
                    # table.font.name = u'Times New Roman'
                    # table._element.rPr.rFonts.set(qn('w:eastAsia'), u'Times New Roman')
                    # table.autofit = True
                    table.cell(0, 0).width = Cm(7.65)
                    table.cell(0, 1).text = series_info[list(series_info.keys())[0]]
                    center_dcms_id = list(center_info[protocal][body_part][series_name].keys())
                    center_dcms_id.sort()

                    margin_dict = {'top': 0, 'start': 0, 'bottom': 0, 'end': 0, 'left': 0, 'right': 0}
                    self.__set_margin(table.cell(0, 0), margin_dict)
                    table_run = table.cell(0, 0).paragraphs[0].add_run()

                    for instance_id in center_dcms_id:
                        pic_path = self.convert_jpg(center_info[protocal][body_part][series_name][instance_id])
                        pic_size = 2.53
                        if len(center_dcms_id) <= 3:
                            pic_width = pic_size * 3 / len(center_dcms_id)
                            table_run.add_picture(pic_path, height=Cm(pic_size * 2), width=Cm(pic_width))     # todo auto fit
                        else:
                            table_run.add_picture(pic_path, height=Cm(pic_size), width=Cm(pic_size))  # todo auto fit
                        os.remove(pic_path)
                docx.add_section()

        docx.save(docx_path)
        print('Write parameters to docx done!')
        print('The output file path in %s' % docx_path)

    def __set_margin(self, cell, margin_dict):
        tc = cell._tc
        tcPr = tc.get_or_add_tcPr()
        tcMar = OxmlElement('w:tcMar')
        for m in ['top', 'start', 'bottom', 'end', 'left', 'right']:
            if m in margin_dict:
                node = OxmlElement('w:{}'.format(m))
                node.set(qn('w:w'), str(margin_dict.get(m)))
                node.set(qn('w:type'), 'dxa')
                tcMar.append(node)
        tcPr.append(tcMar)

    def extract_info_2_docx(self):
        tag_info, center_info = self.extract_parameter()
        self.write_2_docx(tag_info, center_info)

if __name__ == '__main__':
    dicom_path = os.path.normpath(filedialog.askdirectory())
    tool_obj = FDATool(dicom_path)
    tool_obj.extract_info_2_docx()
    os.startfile(r'.\Extracted')